package com.maiko.maikoaiagent.demo.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.stereotype.Component;

import java.util.List;
/**
 * MultiQueryExpanderDemo 类用于演示如何使用多查询扩展功能
 * 它通过与ChatClient的交互，将一个查询扩展为多个相关的查询
 *
 * @author: Maiko7
 * @create: 2025-07-09 10:21
 */
@Component
public class MultiQueryExpanderDemo {

    // ChatClient.Builder 的实例，用于构建ChatClient
    private final ChatClient.Builder chatClientBuilder;

    /**
     * 构造方法
     * 初始化MultiQueryExpanderDemo类的实例，需要一个ChatModel类型的参数
     *
     * @param dashscopeChatModel 用于构建ChatClient的ChatModel实例
     */
    public MultiQueryExpanderDemo(ChatModel dashscopeChatModel) {
        this.chatClientBuilder = ChatClient.builder(dashscopeChatModel);
    }

    /**
     * 扩展查询
     * 将给定的查询字符串扩展为多个查询，返回一个包含多个Query对象的列表
     *
     * @param query 需要扩展的原始查询字符串
     * @return 包含多个扩展后的查询对象的列表
     */
    public List<Query> expand(String  query) {
        // 构建MultiQueryExpander实例，指定聊天客户端构建器和查询数量
        MultiQueryExpander queryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(chatClientBuilder)       // 传入构建好的 ChatClient.Builder
                .numberOfQueries(3)      // 设置生成查询的数量为 3 个
                .build();

        // 使用MultiQueryExpander扩展原始查询，并返回扩展后的查询列表
        List<Query> queries = queryExpander.expand(new Query("谁是程序员Maiko啊？"));
        return queries;
    }

}

